#include "wo_routine.h"
#include "wo_types.h"
#include <assert.h>
#include <stdbool.h>
#include <stddef.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <threads.h>

// coroutine scheduler instance. thread local variable.
_Thread_local wo_scheduler_t *S;

// swap context. save current call stack into swap_out then restore swap_in.
void WOAPI wo_swap_context(wo_context_t *swap_out,
                           wo_context_t *const swap_in) {
  UNUSED(swap_out);
  UNUSED(swap_in);
#if defined(__x86_64__)
  __asm__("movq $0, %rax\n\t"
          "movq %rbx, 0x8(%rdi)\n\t"
          "movq %rcx, 0x10(%rdi)\n\t"
          "movq %rdx, 0x18(%rdi)\n\t"
          "movq %rsi, 0x20(%rdi)\n\t"
          "movq %rdi, 0x28(%rdi)\n\t"
          "movq %rbp, 0x30(%rdi)\n\t"
          "movq %rsp, 0x38(%rdi)\n\t"
          "movq %r8,  0x40(%rdi)\n\t"
          "movq %r9,  0x48(%rdi)\n\t"
          "movq %r10, 0x50(%rdi)\n\t"
          "movq %r11, 0x58(%rdi)\n\t"
          "movq %r12, 0x60(%rdi)\n\t"
          "movq %r13, 0x68(%rdi)\n\t"
          "movq %r14, 0x70(%rdi)\n\t"
          "movq %r15, 0x78(%rdi)\n\t"
          "\n\t"
          "movq 0x8(%rsi), %rbx\n\t"
          "movq 0x10(%rsi), %rcx\n\t"
          "movq 0x18(%rsi), %rdx\n\t"
          "movq 0x20(%rsi), %rdi\n\t"
          "movq 0x30(%rsi), %rbp\n\t"
          "movq 0x38(%rsi), %rsp\n\t"
          "movq 0x40(%rsi), %r8\n\t"
          "movq 0x48(%rsi), %r9\n\t"
          "movq 0x50(%rsi), %r10\n\t"
          "movq 0x58(%rsi), %r11\n\t"
          "movq 0x60(%rsi), %r12\n\t"
          "movq 0x68(%rsi), %r13\n\t"
          "movq 0x70(%rsi), %r14\n\t"
          "movq 0x78(%rsi), %r15\n\t"
          // 最后覆盖 rsi 寄存器，防止中途把参数改了
          "movq 28(%rsi), %rsi\n\t");
#elif defined(__i386__)
  // old ebp     <- ebp
  // return addr <- ebp + 4
  // swap_out    <- ebp + 8
  // swap_in     <- ebp + 12 (0xc)
  __asm__("movl 0x8(%ebp), %eax\n\t" // 备份寄存器状态
          "movl %ebx, 0x4(%eax)\n\t"
          "movl %ecx, 0x8(%eax)\n\t"
          "movl %edx, 0xc(%eax)\n\t"
          "movl %esi, 0x10(%eax)\n\t"
          "movl %edi, 0x14(%eax)\n\t"
          "movl %ebp, 0x18(%eax)\n\t"
          "movl %esp, 0x1c(%eax)\n\t"
          "\n\t"
          "movl 0xc(%ebp),  %eax\n\t" // 恢复寄存器状态
          "movl 0x4(%eax),  %ebx\n\t"
          "movl 0x8(%eax),  %ecx\n\t"
          "movl 0xc(%eax),  %edx\n\t"
          "movl 0x10(%eax), %esi\n\t"
          "movl 0x14(%eax), %edi\n\t"
          "movl 0x18(%eax), %ebp\n\t"
          "movl 0x1c(%eax), %esp\n\t");
#endif
  return;
}

// 协程退出
void WOAPI wo_exit(void) {
  S->C->state = WO_STATE_EXITED;
  wo_swap_context(&S->C->ctx, &S->ctx);
  return;
}

// 保存旧上下文，加载新上下文的 rsp 寄存器，然后在新上下文调用指定函数
void WOAPI wo_start_context(wo_context_t *swap_out, wo_context_t *const swap_in,
                            wo_routine_func_t func, void *args) {
  UNUSED(swap_out);
  UNUSED(swap_in);
  UNUSED(func);
  UNUSED(args);
#if defined(__x86_64__)
  __asm__("movq %rbx, 0x8(%rdi)      \n\t"
          "movq %rcx, 0x10(%rdi)     \n\t"
          "movq %rdx, 0x18(%rdi)     \n\t"
          "movq %rsi, 0x20(%rdi)     \n\t"
          "movq %rdi, 0x28(%rdi)     \n\t"
          "movq %rbp, 0x30(%rdi)     \n\t"
          "movq %rsp, 0x38(%rdi)     \n\t"
          "movq %r8,  0x40(%rdi)     \n\t"
          "movq %r9,  0x48(%rdi)     \n\t"
          "movq %r10, 0x50(%rdi)     \n\t"
          "movq %r11, 0x58(%rdi)     \n\t"
          "movq %r12, 0x60(%rdi)     \n\t"
          "movq %r13, 0x68(%rdi)     \n\t"
          "movq %r14, 0x70(%rdi)     \n\t"
          "movq %r15, 0x78(%rdi)     \n\t"
          "                          \n\t"
          "movq 0x38(%rsi), %rsp     \n\t"
          "movq %rcx,%rdi            \n\t"
          "call *%rdx                \n\t"
          "call wo_exit              \n\t");
#elif defined(__i386__)
  // cdecl 约定
  // old ebp <- ebp
  // return addr <- ebp + 4
  // swap_out <- ebp + 8
  // swap_in <- ebp + 12 (0xc)
  // func <- ebp + 16 (0x10)
  // args <- ebp + 20 (0x14)
  __asm__("movl 0x8(%ebp), %eax     \n\t" // 保存寄存器
          "movl %ebx, 0x4(%eax)     \n\t"
          "movl %ecx, 0x8(%eax)     \n\t"
          "movl %edx, 0xc(%eax)     \n\t"
          "movl %esi, 0x10(%eax)    \n\t"
          "movl %edi, 0x14(%eax)    \n\t"
          "movl %ebp, 0x18(%eax)    \n\t"
          "movl %esp, 0x1c(%eax)    \n\t"
          "                         \n\t"
          "movl 0xc(%ebp), %eax     \n\t" // 设置协程栈
          "movl 0x1c(%eax), %esp    \n\t"
          "                         \n\t"
          "movl 0x10(%ebp), %eax    \n\t" // 调用协程函数
          "pushl 0x14(%ebp)         \n\t"
          "call *%eax               \n\t"
          "addl $0x4, %esp          \n\t");
  wo_exit();
#endif
}

// 从工作协程恢复到调度协程继续执行
void WOAPI wo_yield(void) {
  S->C->state = WO_STATE_SUSPENDED;
  wo_swap_context(&S->C->ctx, &S->ctx);
  return;
}

// 保存协程状态，转移到调度器上下文
void WOAPI wo_resume(void) {
  if (S->C->state == WO_STATE_READY) {
    S->C->state = WO_STATE_RUNNING;
    wo_start_context(&S->ctx, &S->C->ctx, S->C->func, S->C->args);
  } else {
    S->C->state = WO_STATE_RUNNING;
    wo_swap_context(&S->ctx, &S->C->ctx);
  }
  return;
}

// 加入协程队列准备运行
void WOAPI wo_start(void (*func)(void *), void *args) {
  if (S == NULL) {
    S = (wo_scheduler_t *)calloc(1, sizeof(wo_scheduler_t));
  }

  wo_routine_t *routine = (wo_routine_t *)calloc(1, sizeof(wo_routine_t));
  routine->stack = (uint8_t *)malloc(STACK_SIZE);
  routine->ctx.sp = (reg_t)(routine->stack + STACK_SIZE);
  memset(routine->stack, 0, STACK_SIZE);
  routine->func = func;
  routine->args = args;
  if (S->C == NULL) {
    S->C = routine;
    return;
  }

  if (S->C->next != NULL) {
    routine->next = S->C->next;
    S->C->next->prev = routine;
  } else {
    routine->next = S->C;
    S->C->prev = routine;
  }
  routine->prev = S->C;
  S->C->next = routine;
}

void WOAPI wo_schedule(void) {
  if (S == NULL || S->C == NULL) {
    return;
  }
  for (;;) {
    if (S->C == NULL) {
      break;
    }

    wo_resume();

    while (S->C->state == WO_STATE_SUSPENDED ||
           S->C->state == WO_STATE_IO_READY) {
      switch (S->C->state) {
      case WO_STATE_EXITED:
        // 从工作队列移除已结束的协程
        if (S->C != S->C->next && S->C->next != NULL) {
          S->C->next->prev = S->C->prev;
          S->C->prev->next = S->C->next;
          S->C = S->C->next;
          break;
        } else {
          return;
        }
      case WO_STATE_SUSPENDED:
        if (S->C != S->C->next && S->C->next != NULL) {
          S->C = S->C->next;
        }
        break;
      default:
        assert(false);
        return;
      }
    }
  }
}

// 协程主函数，记录栈底地址并启动调度协程
void WOAPI wo_main(void) {
  if (S == NULL) {
    return;
  }

  wo_schedule();
}
